`include "Header.svh"

module Tensor_Core #(
    parameter DATA_WIDTH = 32
  )(
    input logic                   clk,
    input logic                   rst_n,
    input logic                   bram1_clk, // 时钟信号
    input logic                   bram1_rst,
    input logic [7:0]             bram1_addr,// 地址总线
    input logic [DATA_WIDTH-1:0]  bram1_din, // 输入数据
    output logic [DATA_WIDTH-1:0] bram1_dout,// 输出数据
    input logic [3:0]             bram1_we, // 写使能(1:写, 0:读)
    input logic                   bram1_en, // 模块使能
    input logic                   bram2_clk,// 时钟信号
    input logic                   bram2_rst,
    input logic [7:0]             bram2_addr,// 地址总线
    input logic [DATA_WIDTH-1:0]  bram2_din, // 输入数据
    output logic [DATA_WIDTH-1:0] bram2_dout,// 输出数据
    input logic [3:0]             bram2_we, // 写使能(1:写, 0:读)
    input logic                   bram2_en, // 模块使能
    input logic                   Start,//Ctrol
    input logic [7:0]             M,//Ctrol
    input logic [7:0]             N,//Ctrol
    input logic [7:0]             K,//Ctrol
    input logic [1:0]             Buffer_Sel,//A/B/C martix
    input logic [1:0]             Cal_Mode,//INT4,INT8,FP16,FP32
    output logic                  done  //Calculaiton Finished
  );
  //Clog2
  // function called clogb2 that returns an integer which has the 
	// value of the ceiling of the log base 2.                      
	function integer clogb2 (input integer bit_depth);              
	  begin                                                           
	  for(clogb2=0; bit_depth>0; clogb2=clogb2+1)                   
	    bit_depth = bit_depth >> 1;                                 
	  end                                                           
	endfunction 
  /****************************************************************************************************

             -----------------------------AXI connect bram Block--------------------------------

  ****************************************************************************************************/
  // Storage array (physical storage units of analog BRAM)
  logic [DATA_WIDTH-1:0] memory_A [0:255];//A Support to 32x16
  logic [DATA_WIDTH-1:0] memory_B [0:255];//B
  logic [DATA_WIDTH-1:0] memory_C [0:255];//C
  logic [DATA_WIDTH-1:0] memory_Result [0:255];
  // Synchronous sequential logic(Slave)
  always_ff @(posedge bram1_clk) begin : Slave_bram
    if (bram1_en) begin
      if (bram1_we[0]) begin
        case (Buffer_Sel)
          2'd0 : memory_A[bram1_addr] <= bram1_din;
          2'd1 : memory_B[bram1_addr] <= bram1_din;
          2'd2 : memory_C[bram1_addr] <= bram1_din;
          default: ;
        endcase
      end else begin
        case (Buffer_Sel)
          2'd0 : bram1_dout <= memory_A[bram1_addr];
          2'd1 : bram1_dout <= memory_B[bram1_addr]; 
          2'd2 : bram1_dout <= memory_C[bram1_addr];
          2'd3 : bram1_dout <= memory_Result[bram1_addr];  
          default: bram1_dout <= 32'hdeadbeef;
        endcase             
      end
    end
  end
  //Simulation Look
    logic [31:0] Result_1[0:15][0:15];
    generate
      for(genvar a1 = 0;a1 < 16;a1=a1+1)begin
        for(genvar a2 = 0;a2<16;a2=a2+1)begin
          assign Result_1[a1][a2] = memory_Result[a1*16+a2];
        end
      end
    endgenerate
  // Synchronous sequential logic(Master)
  //Read the data directly through combinational logic; otherwise, it will be misaligned by one bit.
    assign bram2_dout = memory_Result[bram2_addr];//Read Only In order to meet the requirements of write operations
  /****************************************************************************************************

             -----------------------------Caluculation Block--------------------------------

  ****************************************************************************************************/
  //Start PE Calculation operation
  //Start Signal Control Tensor Core start Work.There is its logic.
  //Data Controler
  //A martix idx: control partitioning of A matrix
    idx_t position_a               ;
    pack_data data_a[0:`M0-1][7:0]  ;
  //B martix idx: control partitioning of B matrix
    idx_t position_b               ;
    pack_data data_b[0:`N0-1][7:0]  ;
  //psum_en
    logic psum_en[0:`M0-1];
  //data_out_valid idx_finish
    logic data_out_valid[0:`M0-1][0:`N0-1];
  //A*B result: martix_mulple(position and data)
    idx_t position_do[0:`M0-1][0:`N0-1]       ;
    logic [31:0] data_result[0:`M0-1][0:`N0-1];
  //Result signed Counter 
    logic [255:0] memory_finish;
    assign done = (memory_finish == {256{1'b1}});
  //Idx partition number
    logic [4:0] M_idx;
    logic [4:0] N_idx;
    logic [4:0] K_idx;
    assign M_idx = M/`M0;
    assign N_idx = N/`N0;
    assign K_idx = K/`K0;
  //Start enable Cube Calculation
    logic [4:0] k_idx_count;
    logic [4:0] m_idx_count;
    logic [4:0] n_idx_count;
    logic readblock_done;
    logic is_last_block;
    assign is_last_block = (k_idx_count == K_idx-1)&&(n_idx_count == N_idx-1)&&(m_idx_count == M_idx-1);
  //Start = 1 Ctrol Cube enable
    logic enable;
    logic Start_1b;
    logic Strat_2b;
    logic start_pluse;
    assign start_pluse = (Start_1b & ~Strat_2b);
    always_ff @( posedge clk ) begin : start_pluse_product
      if(!rst_n)begin
        Start_1b <= 'd0;
        Strat_2b <= 'd0;
        enable <= 'd0;
      end else begin
        Start_1b <= Start;
        Strat_2b <= Start_1b;
        enable <= (Start & ~readblock_done & ~is_last_block);
      end
    end
  //Ctrol idx_martix 2 Cube, Ctrol enable and Block 
    always_ff @( posedge clk ) begin : mnk2Cube
      if(!rst_n)begin
        k_idx_count <= 'd0;
        m_idx_count <= 'd0;
        n_idx_count <= 'd0;
        readblock_done <= 1'b0;
        for(integer i = 0;i<`M0;i=i+1)begin
          psum_en[i] <= 'd0;
        end
      end else if(enable) begin
        if(m_idx_count < M_idx)begin
          if(n_idx_count == N_idx -1)begin
            if(k_idx_count == K_idx -1)begin
              m_idx_count <= m_idx_count + 'd1;
              n_idx_count <= 'd0;
              k_idx_count <= 'd0;
            end
            else begin
              k_idx_count <= k_idx_count + 'd1;
            end   
          end
        else begin
            if(k_idx_count == K_idx-1)begin
              n_idx_count <= n_idx_count + 'd1;
              k_idx_count <= 'd0;
            end
            else begin
              k_idx_count <= k_idx_count + 'd1;
            end
          end
        end
        for(integer i = 0;i<`M0;i=i+1)begin
          psum_en[i] <= 'd1;
        end
        if(is_last_block)begin
          readblock_done <= 1'b1;
        end
      end
    end
  //Get data and position from A and B martix
    //cube_enable.In order to read the data from A and B martix
      logic cube_enable;
    always_ff @( posedge clk ) begin : Cube_en
      if(!rst_n)begin
        cube_enable <= 'd0;
      end else begin
        cube_enable <= enable;
      end
    end
    //data_a position_a data_b position_b
    always_ff @( posedge clk ) begin : data_information
      if(!rst_n)begin
        position_a <= 'd0;
        position_b <= 'd0;
        for(integer i = 0; i <`M0; i = i+1)begin
          for(integer j=0;j<8;j=j+1)begin
            data_a[i][j] <= 'd0;
          end
        end
        for(integer i =0; i < `N0; i = i+1)begin
          for(integer j=0;j<8;j=j+1)begin
            data_b[i][j] <= 'd0;
          end
        end
      end else if(enable)begin
        position_a.m1_idx <= m_idx_count;
        position_a.k1_idx <= k_idx_count;
        position_b.n1_idx <= n_idx_count;
        position_b.k1_idx <= k_idx_count;
        //Control data_a(Get the m0k0 datas) data_b(Get the k0n0 datas)
          //A m k Ctrol | B k n Ctrol
        //A martix
        for(integer i = 0; i <`M0; i = i+1)begin
          for(integer j=0;j<8;j=j+1)begin
            data_a[i][j] <= memory_A[(m_idx_count*`M0 + i)*K + k_idx_count*`K0+j];
          end
        end
        //B martix
        for(integer i =0; i < `N0; i = i+1)begin
          for(integer j=0;j<8;j=j+1)begin
            data_b[i][j] <= memory_B[(k_idx_count*`K0 + j)*N + n_idx_count*`N0 + i];
          end
        end
      end 
    end
  
  //Cube instantiation
    Cube dut(
      .clk     (clk)   ,
      .rst_n   (rst_n)   ,
      .Start   (cube_enable)   ,
      .Cal_Mode(Cal_Mode)   ,
      .K_idx   (K_idx)   ,
      .position_a (position_a ),         
      .data_a     (data_a     ), 
      .position_b (position_b ),         
      .data_b     (data_b     ),
      .psum_en(psum_en),
      .position_do(position_do), 
      .data_result(data_result),
      .data_valid(data_out_valid) 
    );
  //Bias Add outside
    //Pipe Line //One-hot encoding //Must do
    idx_t position_reg[`M0][`N0][1:0];
    logic valid_reg   [`M0][`N0][1:0];  
    always_ff @( posedge clk ) begin : pipe_line
      if(!rst_n)begin
        for(integer m = 0;m<`M0;m=m+1)begin
          for(integer n = 0;n<`N0;n=n+1)begin
            for(integer i =0;i<2;i=i+1)begin
              position_reg[m][n][i] <= 'd0;
              valid_reg[m][n][i] <= 'd0;
            end
          end
        end
      end else begin
        for(integer m = 0;m<`M0;m=m+1)begin
          for(integer n=0;n<`N0;n=n+1)begin
            position_reg[m][n][0] <= position_do[m][n];
            valid_reg[m][n][0] <= data_out_valid[m][n];
          end
        end
        for(integer m = 0;m<`M0;m=m+1)begin
          for(integer n=0;n<`N0;n=n+1)begin
            for(integer i =1;i<2;i=i+1)begin
              position_reg[m][n][i] <= position_reg[m][n][i-1];
              valid_reg[m][n][i] <= valid_reg[m][n][i-1];
            end
          end
        end
      end
    end
    //instantiation
    logic signed [15:0] int_bias[`M0][`N0];
    generate
      for(genvar m=0;m<`M0;m=m+1)begin
        for(genvar n=0;n<`N0;n=n+1)begin
          int33_addition bias_add_inst(
            .clk    (clk) ,
            .rst_n  (rst_n) ,
            .calc_en(data_out_valid[m][n]) ,
            .a      (data_result[m][n][15:0]) ,
            .b      (memory_C[(position_do[m][n].m1_idx*`M0 + m)*N + position_do[m][n].n1_idx*`N0 + n]) ,
            .c      (int_bias[m][n]) ,
            .done   ()  
          );
        end
      end
    endgenerate
  //Sum2Bram_result
    //MY_Tensor_Unit(m+m1_idx_done_1b[m][n]*`M0)*COL + n1_idx_done_1b[m][n]*`N0 + n )
    always_ff @( posedge clk ) begin : sum2ram
      for(integer m =0 ; m < `M0;m=m+1)begin
        for(integer n = 0;n<`N0;n=n+1)begin
          if(valid_reg[m][n][0])begin
              memory_Result[(m+position_reg[m][n][0].m1_idx*`M0) * N + position_reg[m][n][0].n1_idx * `N0 + n] <= (int_bias[m][n][0] ? {{16{1'b1}},int_bias[m][n]} : {16'd0,int_bias[m][n]} );
              memory_finish[(m+position_reg[m][n][0].m1_idx*`M0) * N + position_reg[m][n][0].n1_idx * `N0 + n] <= 'd1;
          end
        end
      end
    end
  // initial Memory
  initial begin
      for (integer i = 0; i < 256; i = i + 1) begin
          memory_finish[i] = 0;
          memory_A[i] = 'd0;
          memory_B[i] = 'd0;
          memory_C[i] = 'd0;
          memory_Result[i] = 0;
      end
  end
endmodule

module Cube (
    input logic clk,
    input logic rst_n,
    input logic Start,
    input logic [1:0] Cal_Mode,
    input logic [4:0] K_idx,
    input idx_t position_a,//position_a important is its row and idx K
    input pack_data data_a[0:`M0-1][7:0],
    input idx_t position_b,//position_b important is its idx K and col
    input pack_data data_b[0:`N0-1][7:0],
    input logic psum_en[0:`M0-1],//psum_en Ctrol
    output idx_t position_do[0:`M0-1][0:`N0-1],
    output logic data_valid[0:`M0-1][0:`N0-1],
    output pack_data data_result[0:`M0-1][0:`N0-1]
    );
    //`M0`K0 is A martix, represent row side
    //`K0`N0 is B martix, represent col side
    pack_data row_data[0:`M0-1][0:`N0-1][7:0];
    pack_data col_data[0:`N0-1][0:`M0-1][7:0];
    idx_t row_position[0:`M0-1][0:`N0-1];
    idx_t col_position[0:`N0-1][0:`M0-1];
    logic row_valid[0:`M0-1][0:`N0-1];
    logic col_valid[0:`N0-1][0:`M0-1];
    logic row_psum_en[0:`M0-1][0:`N0-1];
    generate
        for (genvar m=0; m < `M0; m=m+1) begin : row_init
            Delay_line #(
                .LEN(m)
            ) delay_line_inst (
                .clk(clk),
                .rst_n(rst_n),
                .data_in(data_a[m]),
                .data_in_valid(Start),
                .data_out(row_data[m][0]),
                .data_out_valid(row_valid[m][0]),
                .position_di(position_a),
                .position_do(row_position[m][0]),
                .psum_en_di(psum_en[m]),
                .psum_en_do(row_psum_en[m][0])
            );
            for (genvar n = 0; n < `N0-1; n= n + 1 ) begin:row_propagate
                Delay_line #(
                    .LEN(1)
                    ) u_row_delay (
                    .clk(clk),
                    .rst_n(rst_n),
                    .data_in(row_data[m][n]),//PE00 -> PE01-> PE02-> PE03-> PE04-> PE05-> PE06-> PE07
                    .data_in_valid(row_valid[m][n]),
                    .data_out(row_data[m][n+1]),    
                    .data_out_valid(row_valid[m][n+1]),
                    .position_di(row_position[m][n]),
                    .position_do(row_position[m][n+1]),
                    .psum_en_di(row_psum_en[m][n]),
                    .psum_en_do(row_psum_en[m][n+1])
                );
            end
        end
    endgenerate
    generate
        for (genvar n=0; n < `N0; n=n+1) begin : col_init
            Delay_line #(
                .LEN(n)
            ) delay_line_inst (
                .clk(clk),
                .rst_n(rst_n),
                .data_in(data_b[n]),
                .data_in_valid(Start),
                .data_out(col_data[n][0]),
                .data_out_valid(col_valid[n][0]),
                .position_di(position_b),
                .position_do(col_position[n][0]),
                .psum_en_di(),
                .psum_en_do()
            );
            for (genvar m=0; m<`M0-1; m=m+1) begin : col_propagate
                Delay_line #(
                    .LEN(1)
                    ) u_col_delay (
                    .clk(clk),
                    .rst_n(rst_n),
                    .data_in(col_data[n][m]),//PE00 -> PE10-> PE20-> PE30-> PE40-> PE50-> PE60-> PE70
                    .data_in_valid(col_valid[n][m]),
                    .data_out(col_data[n][m+1]),    
                    .data_out_valid(col_valid[n][m+1]),
                    .position_di(col_position[n][m]),
                    .position_do(col_position[n][m+1]),
                    .psum_en_di(),
                    .psum_en_do()
                );
            end
        end
    endgenerate
    //get the right positon. 
    //for Example(martix multiple == row x col)
    //ep1: 0 row from A martix multiple 0 col from B martix, result position is 00
    //ep1: 0 row from A martix multiple 1 col from B maritx, result position is 01
    //ep2: 1 row from A martix multiple 1 col from B martix, result position is 11
    idx_t position_mnk_di [0:`M0-1][0:`N0-1];
    generate
        for(genvar m = 0;m<`M0;m=m+1)begin
            for(genvar n = 0;n<`M0;n=n+1)begin
                assign position_mnk_di[m][n].m1_idx = row_position[m][n].m1_idx;
                assign position_mnk_di[m][n].n1_idx = col_position[n][m].n1_idx;
                assign position_mnk_di[m][n].k1_idx = row_position[m][n].k1_idx;
            end
        end
    endgenerate
    //PE instantiation
    // My question is when Bias_en[m][n] is set to get the right bias_add
    generate
        for (genvar m=0; m<`M0; m=m+1) begin : pe_row
            for (genvar n=0; n<`N0; n=n+1) begin : pe_col
                PE Pe_inst(
                    .clk          (clk),
                    .rst_n        (rst_n),
                    .data_in_valid(row_valid[m][n] & col_valid[n][m]),
                    .Cal_Mode     (Cal_Mode),
                    .K_idx        (K_idx),
                    .a            (row_data[m][n]),
                    .b            (col_data[n][m]),
                    .psum_en      (row_psum_en[m][n]),
                    .position_di  (position_mnk_di[m][n]),
                    .position_do  (position_do[m][n]),
                    .pe_result    (data_result[m][n]),
                    .pe_valid(data_valid[m][n])
                );
            end
        end
    endgenerate
endmodule
